#!/usr/bin/env python3

# (c) Copyright 2023 CORSIKA Project, corsika-project@lists.kit.edu
#
# See file AUTHORS for a list of contributors.
#
# This software is distributed under the terms of the 3-clause BSD license.
# See file LICENSE for a full version of the license.


import pickle, sys, itertools


def load_particledb(filename):
    '''
    loads the pickled particle_db (which is an OrderedDict)
    '''
    with open(filename, "rb") as f:
        particle_db = pickle.load(f)
    return particle_db


def read_fluka_codes(filename, particle_db):
    '''
    reads to FLUKA codes data file

    For particles that can interact in FLUKA, add can-interact flag
    '''
    with open(filename) as f:
        for line in f:
            line = line.strip()
            if len(line) == 0 or line[0].startswith('#'):
                continue
            line = line.split()
            if len(line) == 2:
                canInteractFlag = "no"
                identifier, fluka_code = line
            else:
                identifier, fluka_code, canInteractFlag = line

            try:
                particle_db[identifier]["fluka_code"] = int(fluka_code)
                particle_db[identifier]["fluka_canInteract"] = (canInteractFlag == "yes")
            except KeyError as e:
                raise Exception("Identifier '{:s}' not found in CORSIKA8 particle_db".format(identifier))


def generate_fluka_enum(particle_db):
    '''
     generates the enum to access epos particles by readable names
    '''
    output = ("enum class FLUKACode : int {\n"
              "  Unknown = 0,\n")
    for identifier, pData in particle_db.items():
        if 'fluka_code' in pData:
            output += "  {:s} = {:d},\n".format(identifier, pData['fluka_code'])
    output += "};\n"
    return output


def generate_corsika2fluka(particle_db):    
    '''
    generates the look-up table to convert corsika codes to epos codes
    '''
    string = "static inline std::array<FLUKACode, {:d}> constexpr corsika2fluka = {{\n".format(len(particle_db))
    for identifier, pData in particle_db.items():
        if pData['isNucleus']: continue
        if 'fluka_code' in pData:
            string += "  FLUKACode::{:s}, \n".format(identifier)
        else:
            string += "  FLUKACode::Unknown, // {:s}\n".format(identifier + ' not implemented in FLUKA')
    string += "};\n"
    return string


def generate_fluka_canInteract(particle_db):    
    '''
    generates the look-up table to convert corsika codes to epos codes
    '''
    string = "static inline std::vector<bool> const flukaCanInteract = {\n"
    for identifier, pData in particle_db.items():
        canInteract = pData.get("fluka_canInteract", False)
        string += "  {:s}, // {:s}\n".format(str(canInteract).lower(), identifier)
    string += "};\n"
    return string


if __name__ == "__main__":
    if len(sys.argv) != 3:
        print("usage: {:s} <particle_db.pkl> <fluka_codes.dat>".format(sys.argv[0]), file=sys.stderr)
        sys.exit(1)
        
    print("code_generator.py for FLUKA")
    
    particle_db = load_particledb(sys.argv[1])
    read_fluka_codes(sys.argv[2], particle_db)
    
    with open("Generated.inc", "w") as f:
        print("// this file is automatically generated\n// edit at your own risk!\n", file=f)
        print(generate_fluka_enum(particle_db), file=f)
        print(generate_corsika2fluka(particle_db), file=f)
        print(generate_fluka_canInteract(particle_db), file=f)
